#!/usr/bin/env python

__version__ = 1.01

from getopt import gnu_getopt, GetoptError
from math import exp, pi, sqrt
from os import mkdir
from os.path import exists
from sys import argv, exit

# Constants
Grav_Const = 9.80665    # Gravitational acceleration
Univeral_Gas_Const = 8.31432    # Universal gas constant
Molar_Mass_of_Air = 0.0289644    # Molar mass of Earth air
bulk = 1.42E5    # Bulk modulus of elasticity for dry air

class RocketSpecs:

    def __init__(self, t_B, T, m_E, m_F, m_P, A):
        # Due to lack of discription, the lables for these variables are just guesses as to their true use based on educated guesses.
        self.t_B = t_B #Time of burn (how long it takes for the fuel to be completely burned off)
        self.T = T #Total time of flight
        self.m_E = m_E #Mass of 
        self.m_F = m_F #Mass of fuel
        self.m_P = m_P #Mass of parachute
        self.A = A #Area of burn (usually assumed to be a cylinder in basic calcualtions)
	################################################################################   
  # BAROMETRIC PRESSURE FORMULA
	# 
	# Reason: Beacause the craft will be going through various levels of pressure due to the makeup of the 
	# atmosphere, this will affect the fuel in the craft as well as the craft itself. If you wish to learn more
	# about the topics discussed, please look at these videos:
	# 
	# Barometric Pressure and Temperature: http://www.youtube.com/watch?v=fqdex5ER6kA
	# Weather and air pressure: http://www.youtube.com/watch?v=I0C4QR0OEH0
	#
	#This is a formula for calculating the barometric pressure at any one moment. 
    #For more information regarding this, please direct your attention to this article: http://en.wikipedia.org/wiki/Barometric_formula
    #The variables have been named appropraitely so that anyone can follow along
	#Here are the full names and units of the variables:
	# Static_Pressure = Static pressure (pascals)
    # Standard_Temp = Standard temperature (K)
    # Standard_Temperature_Lapse_Rate = Standard temperature lapse rate -0.0065 (K/m) in ISA
    # Height_Above_Sea = Height above sea level (meters)
    # Height_Above_Sea_Lower = Height at bottom of layer b (meters; e.g.,  = 11,000 meters)
    # Univeral_Gas_Const = Universal gas constant for air: 8.31432 N·m /(mol·K)
    # Grav_Const = Gravitational acceleration (9.80665 m/s2)
    # Molar_Mass_of_Air = Molar mass of Earth's air (0.0289644 kg/mol)		
    ################################################################################
def barometric_formula(Height_Above_Sea, Static_Pressure, Standard_Temperature_Lapse_Rate, Standard_Temp, Height_Above_Sea_Lower):
    if Standard_Temperature_Lapse_Rate == 0:
        return Static_Pressure*exp((-1*Grav_Const*Molar_Mass_of_Air*(Height_Above_Sea-Height_Above_Sea_Lower))/(Univeral_Gas_Const*Standard_Temp))
    else:
        return Static_Pressure*(Standard_Temp/(Standard_Temp+Standard_Temperature_Lapse_Rate*(Height_Above_Sea-Height_Above_Sea_Lower)))**((Grav_Const*Molar_Mass_of_Air)/(Univeral_Gas_Const*Standard_Temperature_Lapse_Rate)+1)
    ################################################################################
	# STATIC PRESSURE VALUE ASSGINMENT
	# 
	# Reason: A lot of values used in the barometric pressure formula are determined through test data that has been
	# compiled by researchers. In this function, a sea level height is used to determine what values should be assgined to 
	# the other variables in the equation. To learn more about air pressure, please watch the videos provided in the BAROMETRIC
	# PRESSURE FORMULA section
	# 
	# Variable List:
    # Static_Pressure = Static pressure (pascals)
    # Standard_Temp = Standard temperature (K)
    # Standard_Temperature_Lapse_Rate = Standard temperature lapse rate -0.0065 (K/m) in ISA
    # Height_Above_Sea = Height above sea level (meters)
    # Height_Above_Sea_Lower = Height at bottom of layer b (meters; e.g.,  = 11,000 meters)
	################################################################################
def Static_Pressure(Height_Above_Sea):
    if 0 <= Height_Above_Sea < 11000:
        Static_Pressure = 1.2250
        Standard_Temperature_Lapse_Rate = -0.0065
        Standard_Temp = 288.15
        Height_Above_Sea_Lower = 0
    elif 11000 <= Height_Above_Sea < 20000:
        Static_Pressure = 0.36391
        Standard_Temperature_Lapse_Rate = 0.0
        Standard_Temp = 216.65
        Height_Above_Sea_Lower = 11000
    elif 20000 <= Height_Above_Sea < 32000:
        Static_Pressure = 0.08803
        Standard_Temperature_Lapse_Rate = 0.001
        Standard_Temp = 216.65
        Height_Above_Sea_Lower = 20000
    elif 32000 <= Height_Above_Sea < 47000:
        Static_Pressure = 0.01322
        Standard_Temperature_Lapse_Rate = 0.0028
        Standard_Temp = 228.65
        Height_Above_Sea_Lower = 32000
    elif 47000 <= Height_Above_Sea < 51000:
        Static_Pressure = 0.00143
        Standard_Temperature_Lapse_Rate = 0.0
        Standard_Temp = 270.65
        Height_Above_Sea_Lower = 47000
    elif 51000 <= Height_Above_Sea < 71000:
        Static_Pressure = 0.00086
        Standard_Temperature_Lapse_Rate = -0.0028
        Standard_Temp = 270.65
        Height_Above_Sea_Lower = 51000
    elif 71000 <= Height_Above_Sea < 86000:
        Static_Pressure = 0.000064
        Standard_Temperature_Lapse_Rate = -0.002
        Standard_Temp = 214.65
        Height_Above_Sea_Lower = 71000
    else:
        return 0.000001

    Static_Pressure = barometric_formula(Height_Above_Sea, Static_Pressure, Standard_Temperature_Lapse_Rate, Standard_Temp, Height_Above_Sea_Lower)
    return Static_Pressure

	

def rk_step(Height_Above_Sea, v, dsdt, dvdt, t, Height_Above_Sea_Lower):
    k1 = dsdt(Height_Above_Sea, v, t)
    k2 = dsdt(Height_Above_Sea+((Height_Above_Sea_Lower*k1)/2), v, t+Height_Above_Sea_Lower/2)
    k3 = dsdt(Height_Above_Sea+((Height_Above_Sea_Lower*k2)/2), v, t+Height_Above_Sea_Lower/2)
    k4 = dsdt(Height_Above_Sea+Height_Above_Sea_Lower*k3, v, t+Height_Above_Sea_Lower)
    Height_Above_Sea += (Height_Above_Sea_Lower/6)*(k1 + 2*k2 + 2*k3 + k4)

    k1 = dvdt(Height_Above_Sea, v, t)
    k2 = dvdt(Height_Above_Sea, v+((Height_Above_Sea_Lower*k1)/2), t+Height_Above_Sea_Lower/2)
    k3 = dvdt(Height_Above_Sea, v+((Height_Above_Sea_Lower*k2)/2), t+Height_Above_Sea_Lower/2)
    k4 = dvdt(Height_Above_Sea, v+Height_Above_Sea_Lower*k3, t+Height_Above_Sea_Lower)
    v += (Height_Above_Sea_Lower/6)*(k1 + 2*k2 + 2*k3 + k4)

    return Height_Above_Sea, v

def compute_stage(dirname, s_0, v_0, dsdt, dvdt, Height_Above_Sea_Lower, t_B):
    t = 0
    Height_Above_Sea = s_0
    v = v_0
    max_a = 0
    if dirname:
        fp_alt = open(dirname + "/altitude","w")
        fp_vel = open(dirname + "/velocity","w")
        fp_mach = open(dirname + "/mach","w")
        fp_acc = open(dirname + "/acceleration","w")
        fp_q = open(dirname + "/pressure","w")
    while t < t_B:
        Height_Above_Sea, v = rk_step(Height_Above_Sea, v, dsdt, dvdt, t, Height_Above_Sea_Lower)
        density = Static_Pressure(Height_Above_Sea)
        speed_of_sound = sqrt(bulk / density)
        mach = v / speed_of_sound
        if dirname:
            fp_alt.write("%f, %f\n" % (t, Height_Above_Sea))
            fp_vel.write("%f, %f\n" % (t, v))
            fp_mach.write("%f, %f\n" % (t, mach))
            fp_acc.write("%f, %f\n" % (t, dvdt(Height_Above_Sea,v,t)))
            fp_q.write("%f, %f\n" % (t, density*v*v/2))
        t += Height_Above_Sea_Lower
    return t, Height_Above_Sea, v

def compute_zenith(dirname, t_0, s_0, v_0, dsdt, dvdt, Height_Above_Sea_Lower):
    t = t_0
    Height_Above_Sea = s_0
    v = v_0
    if dirname:
        fp_alt = open(dirname + "/altitude","a")
        fp_vel = open(dirname + "/velocity","a")
        fp_mach = open(dirname + "/mach","a")
        fp_acc = open(dirname + "/acceleration","a")
        fp_q = open(dirname + "/pressure","a")
    while v > 0:
        Height_Above_Sea, v = rk_step(Height_Above_Sea, v, dsdt, dvdt, t, Height_Above_Sea_Lower)
        density = Static_Pressure(Height_Above_Sea)
        speed_of_sound = sqrt(bulk / density)
        mach = v / speed_of_sound
        if dirname:
            fp_alt.write("%f, %f\n" % (t, Height_Above_Sea))
            fp_vel.write("%f, %f\n" % (t, v))
            fp_mach.write("%f, %f\n" % (t, mach))
            fp_acc.write("%f, %f\n" % (t, dvdt(Height_Above_Sea,v,t)))
            fp_q.write("%f, %f\n" % (t, density*v*v/2))
        t += Height_Above_Sea_Lower
    return Height_Above_Sea, v

def get_drag(Height_Above_Sea, v):
    density = Static_Pressure(Height_Above_Sea)
    speed_of_sound = sqrt(bulk / density)
    mach = v / speed_of_sound
    if mach <= 0.9 or mach > (0.41/0.21):
        drag_coeff = 0.15
    elif 0.9 < mach <= 1.0:
        drag_coeff = 0.15 + (0.26)*(mach-0.9)/0.1
    elif 1.0 < mach <= (0.41/0.21):
        drag_coeff = 0.41 - 0.21*(mach-1.0)
    return drag_coeff

def simulate_flight(dirname, specs):

    if dirname and not exists(dirname):
        mkdir(dirname)

    # Simulate engine burn
    def dsdt(Height_Above_Sea, v, t):
        return v

    def dvdt(Height_Above_Sea, v, t):
        mass = (specs.m_E + specs.m_P + (1-t/specs.t_B)*specs.m_F)
        drag = get_drag(Height_Above_Sea, v)
        return -1*Grav_Const + (1/mass)*(specs.T - 0.5*Static_Pressure(Height_Above_Sea)*drag*specs.A*v*v)

    t, Height_Above_Sea, v = compute_stage(dirname, 0, 0, dsdt, dvdt, 0.1, specs.t_B)

    # Ascent to zenith
    def dvdt(Height_Above_Sea, v, t):
        drag = get_drag(Height_Above_Sea, v)
        return -1*Grav_Const - (0.5*Static_Pressure(Height_Above_Sea)*drag*specs.A*v*v)/(specs.m_E + specs.m_P)

    Height_Above_Sea, v = compute_zenith(dirname, t, Height_Above_Sea, v, dsdt, dvdt, 0.1)

    return Height_Above_Sea

def search(filename, specs, min_thrust=1000, max_thrust=50000):
    min_req_thrust = (specs.m_E+specs.m_F+specs.m_P)*9.8
    specs.t_B = 5
    specs.T = min_thrust
    fp = open("usofs_search", "w")
    while specs.T <= max_thrust:
        if specs.T < min_req_thrust:
            specs.T += 500
            continue
        print "Working on thrust T=%d N..." % specs.T
        specs.t_B = 5
        zenith = 0
        while zenith < 100000:
            zenith = simulate_flight(None, specs)
            specs.t_B += 1
        fp.write("%f, %d\n" % (specs.T, specs.t_B))
        specs.T += 500
    fp.close()

def usage():

    useagestr = """Foo"""
    print usagestr

def version():

    versionstr = """USOFS Version: %s""" % str(__version__)
    print versionstr

def main():

    argstring = "a:b:e:f:ho:p:st:v"
    try:
        opts, args = gnu_getopt(argv[1:], argstring)
    except getopt.GetoptError, err:
        print str(err)
        usage()
        exit(2)
    do_search = False
    output = True
    specs = RocketSpecs(60, 1000, 30, 100, 5, pi*0.15*0.15)
    filename = "usofs_output"
    for o, a in opts:
        if o == "-a":
            specs.A = pi*a*a
        if o == "-b":
            specs.t_B = a
        if o == "-e":
            specs.m_E = a
        if o == "-f":
            specs.m_F = a
        if o == "-h":
            usage()
            exit()
        if o == "-o":
            filename = a
        if o == "-p":
            specs.m_P = a
        if o == "-s":
            do_search = True
            output = False
            filename = "usofs_search"
        if o == "-t":
            specs.T = a
        if o == "-v":
            version()
            exit()

    if do_search:
        search(filename, specs)
        exit()
    else:
        simulate_flight(filename, specs)
        exit()

if __name__ == "__main__":
    main()